import os
from rapidocr_paddle import RapidOCR
from langchain.document_loaders.unstructured import UnstructuredFileLoader
from typing import List
from PyCmpltrtok.common import sep

global_ocr_obj = None


def get_ocr():
    global global_ocr_obj
    if global_ocr_obj is None:
        use_cuda = True
        print(f'**** Loading RapidOCR from rapidocr_paddle with cuda = {use_cuda}')
        global_ocr_obj = RapidOCR(det_use_cuda=use_cuda, cls_use_cuda=use_cuda, rec_use_cuda=use_cuda)
    return global_ocr_obj


class RapidOCRLoader(UnstructuredFileLoader):
    def _get_elements(self) -> List:
        def img2text(filepath):
            resp = ""
            ocr = get_ocr()
            result, _ = ocr(filepath)
            if result:
                ocr_result = [line[1] for line in result]
                resp += "\n".join(ocr_result)
            return resp

        text = img2text(self.file_path)
        from unstructured.partition.text import partition_text
        return partition_text(text=text, **self.unstructured_kwargs)
    
    
if '__main__' == __name__:
    
    import argparse
    
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('names', help='name of the input images', nargs='+')
    parser.add_argument('--dir', help='base dir', type=str, default='.')
    
    args = parser.parse_args()
    names = args.names
    xdir = args.dir
    
    sep('OCR doing')
    for i, name in enumerate(names):
        n = i + 1
        sep(f'{n}')
        path = os.path.join(xdir, name)
        print(f'path={path}')
        loader = RapidOCRLoader(file_path=path)
        docs = loader.load()
        text = '\n\n'.join([doc.page_content for doc in docs])
        print(text)
    